package com.firesoon;

import lombok.extern.slf4j.Slf4j;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.implementation.MethodDelegation;
import net.bytebuddy.matcher.ElementMatchers;

import java.lang.instrument.Instrumentation;

/**
 * @author yakun.zeng
 * @version 1.0
 * @description: TODO
 * @date 2025/7/28 20:18
 */
@Slf4j
public class AgentByPremain {

    private static final String ANNOTATION_PREFIX = "org.springframework.web.bind.annotation";
    private static final String ANNOTATION_CONTROLLER = "org.springframework.stereotype.Controller";
    private static final String ANNOTATION_REST_CONTROLLER = "org.springframework.web.bind.annotation.RestController";

    public static void premain(String agentArgs, Instrumentation inst) {

        log.info("Loading " + AgentByPremain.class.getName() + " ...");

        try {
            // 拦截spring controller
            AgentBuilder.Identified.Extendable agentBuilder = new AgentBuilder.Default()
                    // 拦截 @Controller 和 @RestController
                    .type(ElementMatchers.isAnnotatedWith(
                            ElementMatchers.named(ANNOTATION_CONTROLLER)
                                    .or(ElementMatchers.named(ANNOTATION_REST_CONTROLLER))))
                    .transform((builder, typeDescription, classLoader, javaModule, protectionDomain) ->
                            // 拦截 @RestMapping 或者 @Get/Post/Put/DeleteMapping
                            builder.method(ElementMatchers.isPublic().and(ElementMatchers.isAnnotatedWith(
                                            ElementMatchers.nameStartsWith(ANNOTATION_PREFIX)
                                                    .and(ElementMatchers.nameEndsWith("Mapping")))))
                                    // 拦截后交给 SpringControllerInterceptor 处理
                                    .intercept(MethodDelegation.to(ControllerInterceptor.class)));
            // 装载到 instrumentation 上
            agentBuilder.installOn(inst);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
